import math
import sets

# Util class to choose the method to split and find the best
# attribute.
class AttributeChooser():
  def __init__(self, method = "information_gain"):
    self.method = method

  # Computing the entropy for the set of examples
  def get_entropy(self, examples):
    labels = [example[1] for example in examples]
    unique_labels = [element for element in sets.Set(labels)]
    counts = [labels.count(label) for label in unique_labels]
    total_cnt = len(labels)
    entropy = 0.0
    # Compute Sigma(-p(x)*log(p(x)))
    for idx in range(0, len(unique_labels)):
      label_prob = float(counts[idx]) / float(total_cnt)
      entropy += -1.0 * label_prob * math.log(label_prob)
    return entropy

  # Get the subset of examples that have the given value for the attribute.
  def get_subsets(self, examples, attribute):
    example_subset = {}
    for example in examples:
      feature_vector = example[0]
      if attribute in feature_vector:
        if feature_vector[attribute] not in example_subset:
          example_subset[feature_vector[attribute]] = []
        example_subset[feature_vector[attribute]].append(example)

    return example_subset.values()

  # Get the intrinsic value which is based on the number of splits and the
  # size of splits
  def get_intrinsic_value(self, example_subsets, total_examples):
    intrinsic_val = 0.0
    for example_subset in example_subsets:
      iv_prob = float(len(example_subset)) / float(total_examples)
      intrinsic_val += -1.0 * iv_prob * math.log(iv_prob)
    return intrinsic_val

  # Compute the information gain based on the attribtutes that we can split.
  def info_gain_attribute(self, examples, attributes, gain_ratio=False):
    max_information_gain = 0.0
    best_attribute = attributes.keys()[0]
    current_entropy = self.get_entropy(examples)
    for attribute in attributes.keys():
      # Get the splits
      example_subsets = self.get_subsets(examples, attribute)

      # Compute entropy
      new_entropy = 0.0
      for example_subset in example_subsets:
        normalize_factor = float(len(example_subset)) / float(len(examples))
        new_entropy += normalize_factor * self.get_entropy(example_subset)
      current_information_gain = current_entropy - new_entropy

      # If using gain ratio divide by intrinsic value
      if gain_ratio:
        intrinsic_val = self.get_intrinsic_value(example_subsets,
            len(examples))
        current_information_gain /= intrinsic_val 

      # Get the max information gain
      if current_information_gain > max_information_gain:
        max_information_gain = current_information_gain
        best_attribute = attribute 

    return best_attribute

  # Choose the best attribute depending on the method (to select the best
  # attribute)
  def choose_best_attribute(self, examples, attributes):
    if self.method == 'information_gain':
      return self.info_gain_attribute(examples, attributes)
    elif self.method == 'gain_ratio':
      return self.info_gain_attribute(examples, attributes, True)
